# 导入相关库
import lightgbm as lgb
import pandas as pd
import numpy as np
cat = {
    'model': 'category',
    'brand': 'category',
    'bodyType': 'category',
    'fuelType': 'category',
    'gearbox': 'category',
    'notRepairedDamage': 'category',
    'regionCode': 'category',
    'seller': 'category',
    'offerType': 'category',
}

#----------导入数据----------
test = pd.read_csv('./data/used_car_testB_20200421.csv', sep=' ', dtype=cat)
train = pd.read_csv('./data/used_car_train_20200313.csv', sep=' ', dtype=cat)
train['train'] = 1
test['train'] = 0
df = pd.concat([train, test], ignore_index=True)
#----------格式转换----------
df['notRepairedDamage'].replace('-', np.nan, inplace=True)
def to_dt(x):
    m = int(x[4:6])
    if m == 0:
        m = 1
    return x[:4] + '-' + str(m) + '-' + x[6:]
df['regDate'] = pd.to_datetime(df['regDate'].astype('str').apply(to_dt))
df['creatDate'] = pd.to_datetime(df['creatDate'].astype('str').apply(to_dt))
df.regDate = df.regDate.astype(int)
df.creatDate = df.creatDate.astype(int)
df['age'] = df.regDate-df.creatDate
'''
for i in ['model', 'bodyType', 'fuelType', 'gearbox', 'notRepairedDamage', 'price']:
    for j in df[df[i].isna()].index:
        try:
            df.loc[j, i] = df[df['name'] == df.loc[j, 'name']][i].mode()[0]
        except:
            pass
'''
#----------删除异常----------
del df['name']
del df['offerType']
del df['seller']
df['power'][df['power'] > 600] = 600
df['power'][df['power'] < 1] = 1
df['v_13'][df['v_13'] > 6] = 6
df['v_14'][df['v_14'] > 4] = 4


#----------填充缺失----------
def predict_na(data, cname):
    train = data.dropna()
    test = data[data[cname].isna()].agg(
        lambda x: x.fillna(x.mode())
    )
    test.drop(cname, axis=1, inplace=True)
    model = lgb.LGBMClassifier(
        boosting_type='gbdt'
    )
    #训练模型
    model.fit(
        train.drop(cname, axis=1), 
        train[cname]
    )
    #对测试集进行预测
    p = model.predict(test)
    for i,j in enumerate(test.index):
        data.loc[j,cname] = p[i]
    return data[cname]


data0 = df.drop(['SaleID', 'price', 'train', 'regionCode'],
                axis=1)  # 剔除功能性特征与label值
lna = data0.isna().sum()[data0.isna().sum() != 0].index
for i in lna:
    p = pd.Categorical(predict_na(data0, i))
    data0[i] = p
    df[i] = p
#----------存储结果----------
df.to_csv('./user_data/df.csv', index=0, sep=' ')
